import pathlib
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from contextlib import contextmanager
import tqdm
from functools import wraps
from queue import Queue
from threading import Event
import queue
import time
import logging
import warnings
import os


class SenseVoiceWrapper:
    def __init__(self,
                 is_recording_event: Event,
                 stt_stopped_event: Event,
                 exit_event: Event,

                 recording_filename_queue: Queue,
                 stt_text_queue: Queue,

                 model_dir: str=str(pathlib.Path.home()) + "/.cache/modelscope/hub/iic",
                 stt_model: str="SenseVoiceSmall",
                 vad_model: str="speech_fsmn_vad_zh-cn-16k-common-pytorch",
                 device: str="cuda:0",
                 show_tqdm: bool=True,
                 show_logs: bool=False,
                 show_warnings: bool=False,
                 ):
        # event
        self.is_recording_event = is_recording_event
        self.stt_stopped_event = stt_stopped_event
        self.exit_event = exit_event
        # queue
        self.recording_filename_queue = recording_filename_queue
        self.stt_text_queue = stt_text_queue
        # 记录是否中断过
        self.was_interrupted = False

        self.model_dir = model_dir
        self.stt_model = model_dir + "/" + stt_model
        self.vad_model = model_dir + "/" + vad_model
        self.device = device

        self.show_tqdm = show_tqdm
        self.show_logs = show_logs
        self.show_warnings = show_warnings

        self._configure_logging()
        self._configure_warnings()

        # 禁用print
        if not self.show_logs:
            import sys
            self.original_stdout = sys.stdout
            sys.stdout = open(os.devnull, 'w')


        self.model = AutoModel(
            model=self.stt_model,
            trust_remote_code=True,
            disable_update=True,
            remote_code="./SenseVoice/model.py",
            vad_model=self.vad_model,
            vad_kwargs={"max_single_segment_time": 30000},
            device=self.device,
        )

        # 恢复print
        if not self.show_logs:
            import sys
            sys.stdout = self.original_stdout

        # 冒充初始化完成标志
        self.stt_stopped_event.set()    
        print("✅ 文字转语音模型 \t| SenseVoice \t→ 加载完成")

    def _configure_logging(self):
        if not self.show_logs:
            # 禁用所有日志
            logging.basicConfig(level=logging.CRITICAL + 1)
            # 特别处理root logger
            logging.getLogger().setLevel(logging.CRITICAL + 1)
            # 处理funasr可能使用的特定logger
            logging.getLogger("root").setLevel(logging.CRITICAL + 1)
    
    def _configure_warnings(self):
        """配置警告输出"""
        if not self.show_warnings:
            # 忽略所有警告
            warnings.filterwarnings("ignore")
        else:
            # 显示所有警告
            warnings.filterwarnings("default")

    @contextmanager
    def _tqdm_controls(self):
        """
        添加一个上下文管理器用于临时控制tqdm
        完全看不懂这这段代码，豆包写的
        """
        if not self.show_tqdm:
            # 保存原始的tqdm构造函数
            original_init = tqdm.tqdm.__init__
            
            # 定义新的构造函数，强制禁用进度条
            @wraps(original_init)
            def new_init(self, *args, **kwargs):
                # 无论传入什么参数，强制设置disable=True
                kwargs['disable'] = True
                original_init(self, *args, **kwargs)
            
            # 替换tqdm的构造函数
            tqdm.tqdm.__init__ = new_init
            
            try:
                yield
            finally:
                # 恢复原始构造函数
                tqdm.tqdm.__init__ = original_init
        else:
            yield

    def transcribe(self, audio_path: str=""):
        """
        转录音频文件
        """
        with self._tqdm_controls():
            # 用于测试的默认音频
            if audio_path == "":
                audio_path = self.model_dir + "/SenseVoiceSmall/example/zh.mp3"

            res = self.model.generate(
                input=audio_path,
                cache={},
                language="auto", 
                use_itn=True,
                batch_size_s=60,
                merge_vad=True, 
                merge_length_s=15,
            )
            text = rich_transcription_postprocess(res[0]["text"])
            return text
        
    def _clear_queue(self, queue: Queue):
        """清空队列中的所有音频文件"""
        while not queue.empty():
            try:
                queue.get_nowait()
                queue.task_done()
            except queue.Empty:
                break
        
    def transcribe_loop(self):
        """主循环"""
        while not self.exit_event.is_set():
            if self.is_recording_event.is_set():
                if not self.was_interrupted:
                    self._clear_queue(self.recording_filename_queue)
                    self.stt_stopped_event.set()
                    self.was_interrupted = True
                time.sleep(0.02)
                continue
            else:
                self.was_interrupted = False
        
            try:
                # 循环等待时间
                filname = self.recording_filename_queue.get(timeout=0.5)

                # 再次检查录音状态
                if self.is_recording_event.is_set():
                    self._clear_queue(self.recording_filename_queue)
                    self.stt_stopped_event.set()
                    self.was_interrupted = True
                    continue

                # 转录
                text = self.transcribe(filname)
                self.stt_text_queue.put(text)
                self.recording_filename_queue.task_done()
            except queue.Empty:
                continue

    def __del__(self):
        if not self.show_logs:
            import sys
            sys.stdout = self.original_stdout

    def start(self):
        self.transcribe_loop()


if __name__ == "__main__":
    pass